#include "cudaHuffman.cuh"

#ifdef WIN32
#include <malloc.h>
#define alloca _alloca
#else
#include <netint/in.h>
#endif

static unsigned long
numbytes_from_numbits(unsigned long numbits)
{
  return numbits / 8 + (numbits % 8 ? 1 : 0);
}

/*
 * get_bit returns the ith bit in the bits array
 * in the 0th position of the return value.
 */
static unsigned char
get_bit(unsigned char* bits, unsigned long i)
{
	return (bits[i / 8] >> i % 8) & 1;
}


__device__ int cuda_huffman_decode_memory(const unsigned char *bufin,
                          unsigned int bufinlen,
                          unsigned char **pbufout,
                          unsigned int *pbufoutlen)
{
  huffman_node *root, *p;
  unsigned int data_count;
  unsigned int i = 0;
  unsigned char *buf = NULL;
  unsigned int bufcur = 0;

  /* Ensure the arguments are valid. */
  if(!pbufout || !pbufoutlen)
    return 1;

  /* Read the Huffman code table. */
  // root = read_code_table_from_memory(bufin, bufinlen, &i, &data_count);
  if(!root)
    return 1;

  //buf = (unsigned char*)malloc(data_count);

  /* Decode the memory. */
  p = root;
  for(; i < bufinlen && data_count > 0; ++i) 
  {
    unsigned char byte = bufin[i];
    unsigned char mask = 1;
    while(data_count > 0 && mask)
    {
      p = byte & mask ? p->one : p->zero;
      mask <<= 1;

      if(p->isLeaf)
      {
        buf[bufcur++] = p->symbol;
        p = root;
        --data_count;
      }
    }
  }

  // free_huffman_tree(root);
  *pbufout = buf;
  *pbufoutlen = bufcur;
  return 0;
}

static huffman_node*
read_code_table_from_memory(const unsigned char* bufin,
							unsigned int bufinlen,
							unsigned int *pindex,
							unsigned int *pDataBytes)
{
	huffman_node *root = new_nonleaf_node(0, NULL, NULL);
	unsigned int count;
	
	/* Read the number of entries.
	   (it is stored in network byte order). */
	if(memread(bufin, bufinlen, pindex, &count, sizeof(count)))
	{
		free_huffman_tree(root);
		return NULL;
	}

	count = ntohl(count);

	/* Read the number of data bytes this encoding represents. */
	if(memread(bufin, bufinlen, pindex, pDataBytes, sizeof(*pDataBytes)))
	{
		free_huffman_tree(root);
		return NULL;
	}

	*pDataBytes = ntohl(*pDataBytes);

	/* Read the entries. */
	while(count-- > 0)
	{
		unsigned int curbit;
		unsigned char symbol;
		unsigned char numbits;
		unsigned char numbytes;
		unsigned char *bytes = NULL;
		huffman_node *p = root;

		if(memread(bufin, bufinlen, pindex, &symbol, sizeof(symbol)))
		{
			free_huffman_tree(root);
			return NULL;
		}

		if(memread(bufin, bufinlen, pindex, &numbits, sizeof(numbits)))
		{
			free_huffman_tree(root);
			return NULL;
		}
		
		numbytes = (unsigned char)numbytes_from_numbits(numbits);
		// bytes = (unsigned char*)malloc(numbytes);
		if(memread(bufin, bufinlen, pindex, bytes, numbytes))
		{
			// free(bytes);
			free_huffman_tree(root);
			return NULL;
		}

		/*
		 * Add the entry to the Huffman tree. The value
		 * of the current bit is used switch between
		 * zero and one child nodes in the tree. New nodes
		 * are added as needed in the tree.
		 */
		for(curbit = 0; curbit < numbits; ++curbit)
		{
			if(get_bit(bytes, curbit))
			{
				if(p->one == NULL)
				{
					p->one = curbit == (unsigned char)(numbits - 1)
						? new_leaf_node(symbol)
						: new_nonleaf_node(0, NULL, NULL);
					p->one->parent = p;
				}
				p = p->one;
			}
			else
			{
				if(p->zero == NULL)
				{
					p->zero = curbit == (unsigned char)(numbits - 1)
						? new_leaf_node(symbol)
						: new_nonleaf_node(0, NULL, NULL);
					p->zero->parent = p;
				}
				p = p->zero;
			}
		}
		
		// free(bytes);
	}

	return root;
}

static void
free_huffman_tree(huffman_node *subtree)
{
  if(subtree == NULL)
    return;

  if(!subtree->isLeaf)
  {
    free_huffman_tree(subtree->zero);
    free_huffman_tree(subtree->one);
  }

  free(subtree);
}

static huffman_node*
new_leaf_node(unsigned char symbol)
{
  huffman_node *p = (huffman_node*)malloc(sizeof(huffman_node));
  p->isLeaf = 1;
  p->symbol = symbol;
  p->count = 0;
  p->parent = 0;
  return p;
}

static huffman_node*
new_nonleaf_node(unsigned long count, huffman_node *zero, huffman_node *one)
{
  huffman_node *p = (huffman_node*)malloc(sizeof(huffman_node));
  p->isLeaf = 0;
  p->count = count;
  p->zero = zero;
  p->one = one;
  p->parent = 0;

  return p;
}

static int
memread(const unsigned char* buf,
        unsigned int buflen,
        unsigned int *pindex,
        void* bufout,
        unsigned int readlen)
{
  assert(buf && pindex && bufout);
  assert(buflen >= *pindex);
  if(buflen < *pindex)
    return 1;
  if(readlen + *pindex >= buflen)
    return 1;
  memcpy(bufout, buf + *pindex, readlen);
  *pindex += readlen;
  return 0;
}

